#ifndef _SLMBUILDER_H_
#define _SLMBUILDER_H_

#include "port.h"
#include "slm.h"

class CSlmDiscounter;

class CSlmBuilder {
public:
  static const int SLM_MAX_R = 16;
  typedef CSlm::TFREQ TFREQ;
  typedef CSlm::TPROB TPROB;
  typedef CSlm::TNode TNode;
  typedef CSlm::TLeaf TLeaf;

public:
  CSlmBuilder()
    :m_nWord(0), nlevel(0), level(NULL), cut(NULL), discounter(NULL),
     nr(NULL), breaker(), m_excludes(), bUseLogPr(0) {}
  ~CSlmBuilder()
    { Close(); }

  void Create(int n);
  void SetNumberOfWord(int nWord) {this->m_nWord = nWord;}
  void SetCut(TFREQ threshold[]);
  void SetDiscounter(CSlmDiscounter* dis[]);
  void SetBreakerIds(int nId, TWORDID brks[]);
  void SetExcludeIds(int nId, TWORDID excludes[]);
  void SetUseLogPr(int bUse)
    { bUseLogPr = bUse;}
  void AddNGram(TWORDID *ngram, TFREQ fr);
  void Build();
  void Write(FILE *out);
  void Close();

  double getPr(int n, TWORDID *w); // get pr(w[n-1]|w[0][]....[n-2])

public:
  typedef std::vector<TNode> TNodeLevel;
  typedef std::vector<TLeaf> TLeafLevel;
  typedef TNodeLevel::iterator TNodeIterator;
  typedef TLeafLevel::iterator TLeafIterator;

protected:
  bool isBreakId(TWORDID id);
  bool isExcludeId(TWORDID id);
  void CountNr();
  void AppendTails();
  void Cut();
  void Discount();
  void CalcBow();
  void *FindChild(int lvl, TNode *root, TWORDID id);
  int CutNodeLevel(TNodeIterator pfirst, TNodeIterator plast,
                   TNodeIterator chfirst, TNodeIterator chlast, int thrd);
  int CutLeafLevel(TNodeIterator pfirst, TNodeIterator plast,
                   TLeafIterator chfirst, TLeafIterator chlast, int thrd);
private:
  int nlevel, bUseLogPr;
  void **level;                 // level[0] is psudo root level,

  int m_nWord;
  TFREQ *cut;                   // cut[1] is not cut threshold for 1-gram
  CSlmDiscounter **discounter;  // discounter[1] is for 1-gram.
  TFREQ (*nr)[SLM_MAX_R];       // nr[1][slm_max_r] is for 1-gram.
  std::vector<TWORDID> breaker;
  std::vector<TWORDID> m_excludes;
};

class CSlmDiscounter {
public:
  // n is array size, nr is TFREQ[n], nr[0] is corpuse size, or sigma r*nr;
  // nr[1] is number of ngram items with freq 1,
  virtual void init(int n, CSlmBuilder::TFREQ *nr) = 0;

  // freq is the ngram frequence, not the conditional pr
  virtual double discount(int freq) = 0;
  virtual const char* getName() = 0;
};

// Good-Turing discount
class CSlmGTDDiscounter : public CSlmDiscounter {
public:
  CSlmGTDDiscounter(int thredhold = 10, double highfreq_discount = 0.95)
    : thres(thredhold), hd(highfreq_discount), dis(NULL) {}
  virtual void init(int n, CSlmBuilder::TFREQ *nr);
  virtual double discount(int freq);
  virtual const char* getName() {return "Good-Turing";}
protected:
  int thres;
  double hd;
  double *dis;
};

class CSlmAbsoluteDiscounter : public CSlmDiscounter {
public:
  CSlmAbsoluteDiscounter(double substract = 0.0) : c(substract) {}
  // c == 0 mean this value should be count according to r[]
  virtual void init(int n, CSlmBuilder::TFREQ *nr);
  virtual double discount(int freq);
  virtual const char* getName() {return "Absolution";}
protected:
  double c;
};

class CSlmLinearDiscounter : public CSlmDiscounter {
public:
  CSlmLinearDiscounter(double shrink = 0.0) : dis(shrink) {}
  // dis == 0 mean this value should be count according to r[]
  virtual void init(int n, CSlmBuilder::TFREQ *nr);
  virtual double discount(int freq);
  virtual const char* getName() {return "Linear";}
protected:
  double dis;
};  

#endif /* _SLMBUILDER_H_ */
